/* Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

        http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
        limitations under the License.
==============================================================================*/

#include "embedding_update_by_address_tiling.h"
#include "register/op_def_registry.h"

namespace {
    constexpr int32_t EMBEDDING_TYPE_FLOAT16 = 2;
    constexpr int32_t EMBEDDING_TYPE_INT32 = 0;
    constexpr int32_t EMBEDDING_TYPE_FLOAT32 = 1;
}

namespace optiling {
    constexpr int32_t BLOCK_DIM = 48;  // 910b一张卡48个vector核
    constexpr int32_t SIZE_OF_HALF = 2;
    constexpr int32_t SIZE_OF_FLOAT_OR_INT = 4;
    constexpr int32_t MIN_BLOCK_SIZE = 32; // ub空间的数据都要按照32对齐
    constexpr uint64_t UB_LIMIT = 175 * 1024;
    constexpr int32_t USR_SIZE = 256;
    constexpr int32_t SYS_WORKSPACE_SIZE = 16 * 1024 * 1024;
    constexpr int32_t PING_PONG_NUM = 1;

    template<typename T>
    static ge::graphStatus CheckPointer(T *pointer, const char *errorMessage)
    {
        if (pointer == nullptr) {
            printf("%s nullptr\n", errorMessage);
            return ge::GRAPH_FAILED;
        }

        return ge::GRAPH_SUCCESS;
    }

    static ge::graphStatus CheckPositiveInt(int64_t value, const char *errorMessage)
    {
        if (value < 0) {
            printf("%s can not be smaller than 0\n", errorMessage);
            return ge::GRAPH_FAILED;
        }
        if (value == 0) {
            printf("%s is 0, warning!!! check whether input is valid\n", errorMessage);
        }
        return ge::GRAPH_SUCCESS;
    }

    static ge::graphStatus TilingFunc(gert::TilingContext *context)
    {
        if (CheckPointer(context, "Update embeddingType context") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        size_t *currentWorkspace = context->GetWorkspaceSizes(1);
        if (CheckPointer(currentWorkspace, "currentWorkspace") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }
        currentWorkspace[0] = SYS_WORKSPACE_SIZE + USR_SIZE;

        auto inputTensor = context->GetInputTensor(0);
        if (CheckPointer(inputTensor, "GetInputTensor inputTensor") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        int64_t inputShape = inputTensor->GetShapeSize();
        if (CheckPositiveInt(inputShape, "inputShape") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        auto inputTensor1 = context->GetInputTensor(1);
        if (CheckPointer(inputTensor1, "GetInputTensor inputTensor1") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        const int64_t inputShapeTmp = (inputShape > 0) ? inputShape : 1;
        int64_t inputDim = inputTensor1->GetShapeSize() / inputShapeTmp;
        if (CheckPositiveInt(inputDim, "inputDim") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        auto attrs = context->GetAttrs();
        if (CheckPointer(attrs, "GetAttrs attrs") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        auto attrPointer = attrs->GetAttrPointer<int64_t>(0);
        if (CheckPointer(attrPointer, "attrPointer") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        int32_t updateType = *(attrPointer);
        if (updateType != 0 && updateType != 1) {
            printf("Attr update_type is not 0 or 1\n");
            return ge::GRAPH_FAILED;
        }
        ge::DataType inputDatatype = inputTensor1->GetDataType();
        int32_t embeddingType;
        if (inputDatatype == ge::DT_FLOAT16) {
            embeddingType = EMBEDDING_TYPE_FLOAT16;
        } else if (inputDatatype == ge::DT_INT32) {
            embeddingType = EMBEDDING_TYPE_INT32;
        } else {
            embeddingType = EMBEDDING_TYPE_FLOAT32;
        }

        int32_t typeSize = SIZE_OF_FLOAT_OR_INT;
        if (embeddingType == EMBEDDING_TYPE_FLOAT16) {
            typeSize = SIZE_OF_HALF;
        }
        int32_t alignNum = MIN_BLOCK_SIZE / typeSize;

        int32_t inputDimAligned = alignNum * ((inputDim - 1 + alignNum) / alignNum);

        // 每个地址需要占用sizeof(int64_t)个字节，singleDataSize表示每个数据的字节数，需要使用2倍的内存空间，因为每次移动都需要复制一份数据
        int32_t occupyAddressBytesNum =
                sizeof(int64_t) + typeSize * inputDimAligned * PING_PONG_NUM * 2;
        // 一轮计算中最多计算多少个addr，由于地址也要搬到ub，所以需要对齐32
        int64_t addrPerLoop = static_cast<int64_t>(
                UB_LIMIT / static_cast<uint64_t>(occupyAddressBytesNum) & (~3U)); // & (~3U)，保证地址数是4的倍数

        if (CheckPositiveInt(addrPerLoop, "addrPerLoop") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }
        TilingData2 tiling;
        tiling.set_ping_pong_num(PING_PONG_NUM);

        tiling.set_update_type(updateType);
        tiling.set_embedding_type(embeddingType);
        tiling.set_update_dim(static_cast<int32_t>(inputDim));
        tiling.set_addr_nums(static_cast<int32_t>(inputShape));
        tiling.set_addr_per_loop(addrPerLoop);
        tiling.set_type_size(typeSize);
        tiling.set_input_dim_aligned(inputDimAligned);

        context->SetBlockDim(BLOCK_DIM);
        if (CheckPointer(context->GetRawTilingData(), "GetRawTilingData RawTilingData") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
        context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());

        return ge::GRAPH_SUCCESS;
    }
}

namespace ge {
    static ge::graphStatus InferShape(gert::InferShapeContext *context)
    {
        return GRAPH_SUCCESS;
    }
    static ge::graphStatus InferDataType(gert::InferDataTypeContext *context)
    {
        if (optiling::CheckPointer(context, "context") != ge::GRAPH_SUCCESS) {
            return ge::GRAPH_FAILED;
        }

        context->SetOutputDataType(0, ge::DataType(DT_FLOAT));
        return GRAPH_SUCCESS;
    }
}

namespace ops {
    class EmbeddingUpdateByAddress : public OpDef {
    public:
        EmbeddingUpdateByAddress(const char *name) : OpDef(name)
        {
            this->Input("address")
                .ParamType(REQUIRED)
                .DataType({ge::DT_INT64, ge::DT_INT64, ge::DT_INT64})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Input("embedding")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Output("y")
                .ParamType(REQUIRED)
                .DataType({ge::DT_FLOAT, ge::DT_FLOAT16, ge::DT_INT32})
                .Format({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND})
                .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND, ge::FORMAT_ND});
            this->Attr("update_type").AttrType(OPTIONAL).Int(0);
            this->SetInferShape(ge::InferShape)
                .SetInferDataType(ge::InferDataType);

            this->AICore()
                .SetTiling(optiling::TilingFunc);

            OpAICoreConfig aicConfig;
            aicConfig.DynamicCompileStaticFlag(true)
                .DynamicFormatFlag(true)
                .DynamicRankSupportFlag(true)
                .DynamicShapeSupportFlag(true)
                .NeedCheckSupportFlag(false)
                .PrecisionReduceFlag(false);

            this->AICore().AddConfig("ascend910b", aicConfig);
            this->AICore().AddConfig("ascend910", aicConfig);
            this->AICore().AddConfig("ascend910_95", aicConfig);
        }
    };
    OP_ADD(EmbeddingUpdateByAddress);
}
